import java.io.File

import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.SparkSession


object ALS {

  def main(args: Array[String]) {

    val sc = new SparkContext(new SparkConf().setMaster("local[*]").setAppName(""))

    val spark = SparkSession
      .builder
      .appName("ALS")
      .getOrCreate()
    import spark.implicits._

    val frame: DataFrame = spark.read.textFile("作业/resources/example/sample_movielens_ratings.txt")
      .map(parseRating)
      .toDF()

    val ratings = frame
    val Array(training, test) = ratings.randomSplit(Array(0.8, 0.2))

    val als = new ALS()
      .setMaxIter(15)
      .setRegParam(0.01)
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")
    val model = als.fit(training)
//    deleteDir(new File("./output/ALS/"))
//    model.save("./output/ALS")

    model.setColdStartStrategy("drop")
    val predictions = model.transform(test)

    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("rating")
      .setPredictionCol("prediction")
    val rmse = evaluator.evaluate(predictions)
    println(s"Root-mean-square error = $rmse")

    val userRecs = model.recommendForAllUsers(10)
    val movieRecs = model.recommendForAllItems(10)

    val users = ratings.select(als.getUserCol).distinct().limit(3)
    val userSubsetRecs = model.recommendForUserSubset(users, 10)
    val movies = ratings.select(als.getItemCol).distinct().limit(3)
    val movieSubSetRecs = model.recommendForItemSubset(movies, 10)

    userRecs.show()
    movieRecs.show()
    userSubsetRecs.show()
    movieSubSetRecs.show()

    spark.stop()
  }

  def parseRating(str: String): Rating = {
    val fields = str.split("::")
    assert(fields.size == 4)
    Rating(fields(0).toInt, fields(1).toInt, fields(2).toFloat, fields(3).toLong)
  }

  def deleteDir(dir: File): Unit = {
    val files = dir.listFiles()
    files.foreach(f => {
      if (f.isDirectory) {
        deleteDir(f)
      } else {
        f.delete()
//        println("delete file " + f.getAbsolutePath)
      }
    })
    dir.delete()
//    println("delete dir " + dir.getAbsolutePath)
  }
}
case class Rating(userId: Int, movieId: Int, rating: Float, timestamp: Long)

